package com.study.dl4jsample.dogsvscats;

import org.datavec.api.io.labels.ParentPathLabelGenerator;
import org.datavec.api.split.FileSplit;
import org.datavec.api.split.InputSplit;
import org.datavec.image.loader.BaseImageLoader;
import org.datavec.image.recordreader.ImageRecordReader;
import org.deeplearning4j.datasets.datavec.RecordReaderDataSetIterator;
import org.deeplearning4j.eval.Evaluation;
import org.deeplearning4j.nn.api.OptimizationAlgorithm;
import org.deeplearning4j.nn.conf.LearningRatePolicy;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.Updater;
import org.deeplearning4j.nn.conf.inputs.InputType;
import org.deeplearning4j.nn.conf.layers.ConvolutionLayer;
import org.deeplearning4j.nn.conf.layers.DenseLayer;
import org.deeplearning4j.nn.conf.layers.OutputLayer;
import org.deeplearning4j.nn.conf.layers.SubsamplingLayer;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.optimize.listeners.ScoreIterationListener;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.dataset.api.preprocessor.DataNormalization;
import org.nd4j.linalg.dataset.api.preprocessor.ImagePreProcessingScaler;
import org.nd4j.linalg.lossfunctions.LossFunctions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.Random;
//卷积神经网络
public class DogsAndCatsCnnClassifier {

    private static final Logger log = LoggerFactory.getLogger(DogsAndCatsCnnClassifier.class);


    private MultiLayerConfiguration conf=null;
    MultiLayerNetwork model =null;
    //输入图片大小
    private Integer width ;
    private Integer height;

    public Integer getWidth() {
        return width;
    }

    public Integer getHeight() {
        return height;
    }

    public DogsAndCatsCnnClassifier setHeight(Integer height) {
        this.height = height;
        return this ;
    }

    public DogsAndCatsCnnClassifier setWidth(Integer width) {
        this.width = width;
        return this ;
    }
    public DogsAndCatsCnnClassifier setSize(Integer width, Integer height) {
        this.width = width;
        this.height = height;

        return this ;
    }

    //初始化神经网络
    public DogsAndCatsCnnClassifier init(){
        final int nChannels = 3; // Number of input channels


        final int iterations = 10; // Number of training iterations
        final int seed = 123; //
        final int outputNum =2;//dogs and cats =2




        MultiLayerConfiguration.Builder builder = new NeuralNetConfiguration.Builder()
                .seed(seed).iterations(iterations) // Training iterations as above

                /*
                    Uncomment the following for learning decay and bias
                 */
                .learningRate(.01)//.biasLearningRate(0.02)

                .optimizationAlgo(OptimizationAlgorithm.STOCHASTIC_GRADIENT_DESCENT)
                .updater(Updater.RMSPROP) //To configure: .updater(new Nesterovs(0.9))
                .list()
             //classifier.add(Conv2D(64, (3, 3), input_shape = (64, 64, 3), activation = 'relu'))
                .layer(0, new ConvolutionLayer.Builder(5, 5)

                        //nIn and nOut specify depth. nIn here is the nChannels and nOut is the number of filters to be applied
                        .nIn(nChannels)
                        .stride(1,1)
                        .nOut(64)
                        .activation(Activation.RELU)
                        .build())
             //classifier.add(MaxPooling2D(pool_size = (2, 2)))
                .layer(1, new SubsamplingLayer.Builder(SubsamplingLayer.PoolingType.MAX) .kernelSize(2,2)
                        .stride(2,2) .build())
            //classifier.add(Conv2D(32, (3, 3), activation = 'relu'))
                .layer(2, new ConvolutionLayer.Builder(3, 3)
                        //nIn and nOut specify depth. nIn here is the nChannels and nOut is the number of filters to be applied

                        .stride(1,1)
                        .nOut(32)
                        .activation(Activation.RELU)
                        .build())
            //classifier.add(MaxPooling2D(pool_size = (2, 2)))
                .layer(3, new SubsamplingLayer.Builder(SubsamplingLayer.PoolingType.MAX) .kernelSize(2,2)
                        .stride(2,2) .build())
                .layer(4, new DenseLayer.Builder().activation(Activation.RELU).nOut(128).build())
                .layer(5, new OutputLayer.Builder(LossFunctions.LossFunction.NEGATIVELOGLIKELIHOOD)
                        .nOut(outputNum)
                        .activation(Activation.SOFTMAX)
                        .build())

                //classifier.add(Flatten())
                .setInputType(InputType.convolutionalFlat(64,64,3)) //See note below
                .backprop(true).pretrain(false)
                ;
        conf =builder.build();
        model =new MultiLayerNetwork(conf);
        model.init();
        model.setListeners(new ScoreIterationListener(10)); //Print score every 10 iterations


        return this ;
    }

    public void printInfo(){
        System.out.println(model.summary());
    }
    public void train(String path){
        final int CHANNELS =3;
        final int N_LABELS = 2;
        final int batchSize = 64; // Test batch size
        int rngseed = 123;
        Random randNumGen = new Random(rngseed);
        InputSplit trainData = new FileSplit(new File(path), BaseImageLoader.ALLOWED_FORMATS, randNumGen);
        ImageRecordReader reader = new ImageRecordReader(width, height, CHANNELS, new ParentPathLabelGenerator());

        RecordReaderDataSetIterator dataset = new RecordReaderDataSetIterator(reader, batchSize, 1, N_LABELS);
        DataNormalization scaler = new ImagePreProcessingScaler(0, 1);
        dataset.setPreProcessor(scaler);
        try {
            reader.initialize(trainData);
        } catch (IOException e) {
            log.error("",e);
        }

        final int nEpochs = 100; // Number of training epochs

        for( int i=0; i<nEpochs; i++ ) {
            model.fit(dataset);
            Evaluation evalActual = model.evaluate(dataset);

            log.info("*** Completed epoch {} ***{}", i,evalActual.stats());

        }
    }


    public static void main(String args[]) {
        DogsAndCatsCnnClassifier cnn =new DogsAndCatsCnnClassifier().setSize(64,64).init();
        cnn.printInfo();
        cnn.train(args[0]);
    }
}
